package 中等.哈希表;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 给你一个下标从 0 开始的整数数组 nums ，以及整数 modulo 和整数 k 。
 * 请你找出并统计数组中 趣味子数组 的数目。
 * 如果 子数组 nums[l..r] 满足下述条件，则称其为 趣味子数组 ：
 * 在范围 [l, r] 内，设 cnt 为满足 nums[i] % modulo == k 的索引 i 的数量。并且 cnt % modulo == k 。
 * 以整数形式表示并返回趣味子数组的数目。
 * 注意：子数组是数组中的一个连续非空的元素序列。
 *
 * @ https://leetcode.cn/problems/count-of-interesting-subarrays/description/
 * @date 2024/11/24
 */
public class 统计趣味子数组的数目_2845 {

    public static void main(String[] args) {

        System.out.println(countInterestingSubarrays(Arrays.asList(3, 2, 4), 2, 1));

    }

    /**
     * 1，在范围 [l, r] 内，设 cnt 为满足 nums[i] % modulo == k 的索引 i 的数量
     * 转换问题， 如果 nums[i] % modulo == k ，可以看成 nums[i] = 1，求 [l, r]
     * 的和就是数量
     * 2，前缀和
     * (prefixSum[r+1] - prefixSum[l]) % m == k
     * 同余定理
     * prefixSum[r+1] % m - prefixSum[l] % m == k
     * prefixSum[r+1] % m -k ==  prefixSum[l] % m
     */
    public static long countInterestingSubarrays(List<Integer> nums, int modulo, int k) {
        int[] prefixSum = new int[nums.size() + 1];

        for (int i = 0; i < nums.size(); i++) {
            prefixSum[i + 1] = prefixSum[i];
            if (nums.get(i) % modulo == k) {
                prefixSum[i + 1] = prefixSum[i] + 1;
            }
        }
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        long ans = 0;

        for (int i = 1; i < prefixSum.length; i++) {
            // 防止 prefixSum[i] - k 是负数
            int l = (prefixSum[i] - k + modulo) % modulo;
            ans += map.getOrDefault(l, 0);
            map.put(prefixSum[i] % modulo, map.getOrDefault(prefixSum[i] % modulo, 0) + 1);

        }
        return ans;
    }


}
